import logging
import time

import numpy as np
import torch
import torch.nn.functional as F

from torch_geometric.graphgym.checkpoint import load_ckpt, save_ckpt, clean_ckpt
from torch_geometric.graphgym.config import cfg
from torch_geometric.graphgym.loss import compute_loss
from torch_geometric.graphgym.register import register_train
from torch_geometric.graphgym.utils.epoch import is_eval_epoch, is_ckpt_epoch

from graphgps.loss.subtoken_prediction_loss import subtoken_cross_entropy
from graphgps.utils import cfg_to_dict, flatten_dict, make_wandb_name

from deepspeed.profiling.flops_profiler import FlopsProfiler
#from torch.autograd import profiler
from torch.profiler import profile, record_function, ProfilerActivity


# 该函数从给定的 batch 中子采样节点。
# min_k: 每个图中至少保留的节点数。
# ratio: 从每个图中随机选择的节点比例。
# 首先通过 torch.unique 找到所有唯一的图批次。然后，对于每个批次，计算子采样后的索引并返回。
def subsample_batch_index(batch, min_k = 1, ratio = 0.1):
    # # 设置随机种子，确保结果可重复
    torch.manual_seed(0)
    # 找到所有唯一的图批次
    unique_batches = torch.unique(batch.batch)
    # Initialize list to store permuted indices
    # 存储采样的索引
    permuted_indices = []
    for batch_index in unique_batches:
        # Extract indices for the current batch
        # 找到当前批次的所有节点索引
        indices_in_batch = (batch.batch == batch_index).nonzero().squeeze()
        # See how many nodes in the graphs
        # And how many left after subsetting
        # 根据比例计算需要采样的节点数量
        k = int(indices_in_batch.size(0)*ratio)
        # If subsetting gives more than 1, do subsetting
        # 确保采样数量大于最小值
        if k > min_k:
            # 随机打乱节点索引
            perm = torch.randperm(indices_in_batch.size(0))
            # 选择前 k 个索引
            idx = perm[:k]
            # 获取对应的原始节点索引
            idx = indices_in_batch[idx]
            # 对索引进行排序
            idx, _ = torch.sort(idx)
        # Otherwise retain entire graph
        else:
            # 如果节点数量少于 min_k，直接使用所有索引
            idx = indices_in_batch
        # # 将当前批次的采样结果添加到列表中
        permuted_indices.append(idx)
    # 将所有采样的索引合并
    idx = torch.cat(permuted_indices)
    # 返回合并后的索引
    return idx


# 该函数计算 ogbn-arxiv 数据集的交叉熵损失
def arxiv_cross_entropy(pred, true, split_idx):
    true = true.squeeze(-1)
    if pred.ndim > 1 and true.ndim == 1:
        pred_score = F.log_softmax(pred[split_idx], dim=-1)
        loss =  F.nll_loss(pred_score, true[split_idx])
    else:
        raise ValueError("In ogbn cross_entropy calculation dimensions did not match.")
    return loss, pred_score

# @profiler.profile
# def profile_mem_forward(model, batch):
#     pred, true = model(batch)
#     return pred, true

# 定义一个训练周期的函数train_epoch，参数包括日志记录器logger、数据加载器loader、模型model、优化器optimizer、学习率调度器scheduler以及批次累积的数量batch_accumulation
def train_epoch(logger, loader, model, optimizer, scheduler, batch_accumulation):
    # flop related
    # 初始化一些用于性能分析的标志，其中if_mem用于内存分析，if_flop用于浮点运算量（FLOPs）分析，if_select用于选择子图的标志
    if_mem = False
    if_flop = False
    if_select = False
    if if_flop:
        prof = FlopsProfiler(model, None)
        #profile_step = 0
        total_flop_s = 0.
        sample_count = 0
        if if_select:
            total_node = 0

    # 将模型设置为训练模式，这会启用某些层（如Dropout和BatchNorm）的训练行为。
    model.train()
    # 清空优化器中的梯度，以便在新的训练步骤中累积梯度。
    optimizer.zero_grad()
    # 记录训练周期开始的时间，用于后续计算每个批次的处理时间。
    time_start = time.time()
    # 开始遍历数据加载器中的每个批次。iter是当前批次的索引，batch是当前批次的数据。
    for iter, batch in enumerate(loader):
        # 如果启用了选择子图的标志if_select，则通过subsample_batch_index函数对当前批次进行子采样，并创建子图。
        if if_select:
            ratio = 1.0
            idx = subsample_batch_index(batch, min_k = 1, ratio = ratio)
            batch = batch.subgraph(idx)
        # flop related
        # 如果启用了浮点运算量分析if_flop，则开始记录浮点运算量的统计。
        if if_flop: # and iter == profile_step:
            prof.start_profile()
        # 将批次的分割标志设置为'train'，以指示这是训练数据。
        batch.split = 'train'
        # 将批次的数据移动到配置中指定的设备上（例如，GPU）。
        batch.to(torch.device(cfg.device))

        # 将批次数据传入模型进行前向传播，获取预测结果pred和真实标签true。
        pred, true = model(batch)
        # 根据当前数据集的名称，计算损失函数。
        # 如果数据集为'ogbg-code2'，则使用subtoken_cross_entropy计算损失和预测分数。
        if cfg.dataset.name == 'ogbg-code2':
            loss, pred_score = subtoken_cross_entropy(pred, true)
            _true = true
            _pred = pred_score
        # 如果数据集为'ogbn-arxiv'，则根据分割索引计算损失和预测分数，同时将真实标签和预测分数从GPU移动到CPU。
        elif cfg.dataset.name == 'ogbn-arxiv':
            split_idx = loader.dataset.split_idx['train'].to(torch.device(cfg.device))
            loss, pred_score = arxiv_cross_entropy(pred, true, split_idx)
            _true = true[split_idx].detach().to('cpu', non_blocking=True)
            _pred = pred_score.detach().to('cpu', non_blocking=True)
        # 如果数据集名称不是前两者，则使用通用的compute_loss函数计算损失和预测分数，并将结果从GPU移动到CPU。
        else:
            loss, pred_score = compute_loss(pred, true)
            _true = true.detach().to('cpu', non_blocking=True)
            _pred = pred_score.detach().to('cpu', non_blocking=True)

        # 如果启用了浮点运算量分析，则停止记录，并获取总的浮点运算量和参数数量，更新总的浮点运算量统计。计算当前批次的节点数量，如果启用了选择子图，则更新总节点数量。
        if if_flop:
            prof.stop_profile()
            flops = prof.get_total_flops()
            flops_s = flops/1000000000.
            total_flop_s+=flops_s
            sample_count+=len(torch.unique(batch.batch))
            params = prof.get_total_params()
            prof.end_profile()
            if if_select:
                total_node += batch.x.size(0)

        # 对损失进行反向传播，计算梯度。
        loss.backward()
        # Parameters update after accumulating gradients for given num. batches.
        # 在累积一定数量的批次后，进行参数更新。如果启用了梯度裁剪，则裁剪梯度，然后通过优化器更新模型参数，并清空梯度。
        # batch_accumulation是配置参数，表示希望在多少个批次（mini-batch）后更新一次模型参数。
        # 如果当前迭代次数加1能被batch_accumulation整除，或者当前迭代是最后一个批次（iter + 1 == len(loader)），那么就执行以下代码。
        if ((iter + 1) % batch_accumulation == 0) or (iter + 1 == len(loader)):
            # 这里检查是否需要对梯度进行裁剪（即限制梯度的范数）。如果配置中启用了clip_grad_norm，那么就会对模型的所有参数执行梯度裁剪，防止梯度爆炸。
            if cfg.optim.clip_grad_norm:
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            # optimizer.step()会根据当前累积的梯度更新模型的参数。这个更新是基于所选的优化算法（如SGD、Adam等）进行的。
            optimizer.step()
            # 在参数更新之后，使用optimizer.zero_grad()将模型参数的梯度清零，以便在下一次前向传播和反向传播时重新计算新的梯度。这样可以避免累积之前的梯度。
            optimizer.zero_grad()
        # 更新日志记录器的状态，包括真实标签、预测值、损失、学习率、每个批次所用的时间、模型参数以及数据集名称。
        logger.update_stats(true=_true,
                            pred=_pred,
                            loss=loss.detach().cpu().item(),
                            lr=scheduler.get_last_lr()[0],
                            time_used=time.time() - time_start,
                            params=cfg.params,
                            dataset_name=cfg.dataset.name)
        # 重置开始时间，为下一个批次的处理计时。
        time_start = time.time()
    # 如果启用了浮点运算量分析，打印平均浮点运算量和参数数量。
    if if_flop:
        print('################ Print flop')
        print(total_flop_s/sample_count, params)
        print('################ End print flop')
    # 如果启用了内存分析，打印最大分配和保留的CUDA内存。
    if if_mem:
        print('################ Print mem')
        print(torch.cuda.max_memory_allocated() / (1024 ** 2))
        print(torch.cuda.max_memory_reserved() / (1024 ** 2))
        #print(prof_mem.key_averages().table(sort_by="self_cuda_memory_usage", row_limit=10))
        print('################ End print mem')
    # 如果启用了选择子图，打印每个批次的平均节点数。
    if if_select:
        print('################ Print avg nodes')
        print(total_node/sample_count)


# eval_epoch() 函数负责在验证或测试阶段对模型进行评估。它会遍历数据加载器中的批次，计算模型的预测结果和损失，并将这些结果记录下来。
# @torch.no_grad()：这个装饰器用于禁用梯度计算，这样在评估时不会存储计算图，可以提高内存使用效率和计算速度。
@torch.no_grad()
def eval_epoch(logger, loader, model, split='val'):
    # 将模型设置为评估模式，以确保模型在推理时的行为是正确的（如禁用 dropout 和 batch normalization）
    model.eval()
    # 记录当前时间，稍后用于计算评估阶段的耗时。
    time_start = time.time()
    # 遍历数据加载器（loader）中的每个批次（batch）。
    for batch in loader:
        # 为当前批次设置标识（训练、验证或测试），方便后续统计。
        batch.split = split
        # 将批次数据移动到配置文件中指定的设备上，确保计算可以在 GPU 上进行（如果可用）。
        batch.to(torch.device(cfg.device))

        # 这里的逻辑根据模型类型的不同（inductive_edge或其他）选择不同的预测方式。如果模型是 inductive_edge 类型，则返回的结果中包含额外的统计信息（extra_stats）。
        if cfg.gnn.head == 'inductive_edge':
            pred, true, extra_stats = model(batch)
        else:
            pred, true = model(batch)
            extra_stats = {}

        # 这段代码根据数据集的名称计算不同的损失。在 ogbg-code2 数据集中，使用 subtoken_cross_entropy 函数计算损失，并将预测结果和真实值保存。
        if cfg.dataset.name == 'ogbg-code2':
            loss, pred_score = subtoken_cross_entropy(pred, true)
            _true = true
            _pred = pred_score
        elif cfg.dataset.name == 'ogbn-arxiv':
            index_split = loader.dataset.split_idx[split].to(torch.device(cfg.device))
            loss, pred_score = arxiv_cross_entropy(pred, true, index_split)
            _true = true[index_split].detach().to('cpu', non_blocking=True)
            _pred = pred_score.detach().to('cpu', non_blocking=True)
        else:
            loss, pred_score = compute_loss(pred, true)
            _true = true.detach().to('cpu', non_blocking=True)
            _pred = pred_score.detach().to('cpu', non_blocking=True)
        # 更新评估统计信息，包括真实标签、预测分数、损失值、耗时等
        logger.update_stats(true=_true,
                            pred=_pred,
                            loss=loss.detach().cpu().item(),
                            lr=0, time_used=time.time() - time_start,
                            params=cfg.params,
                            dataset_name=cfg.dataset.name,
                            **extra_stats)
        # 最后，重新记录当前时间，为下一个批次的评估做准备。
        time_start = time.time()


# @register_train('custom')：将 custom_train 函数注册为训练方法之一，以便在其他地方调用。
@register_train('custom')
def custom_train(loggers, loaders, model, optimizer, scheduler):
    """
    Customized training pipeline.

    Args:
        loggers: List of loggers
        loaders: List of loaders
        model: GNN model
        optimizer: PyTorch optimizer
        scheduler: PyTorch learning rate scheduler

    """
    # 初始化起始轮数 start_epoch 为0
    start_epoch = 0

    # 如果配置中启用了自动恢复 (cfg.train.auto_resume)，则调用 load_ckpt() 函数加载最近的检查点并设置起始轮数。
    if cfg.train.auto_resume:
        start_epoch = load_ckpt(model, optimizer, scheduler,
                                cfg.train.epoch_resume)

    # 如果起始轮数等于最大轮数，记录日志表示任务已完成；否则，记录从哪个轮次开始训练。
    if start_epoch == cfg.optim.max_epoch:
        logging.info('Checkpoint found, Task already done')
    else:
        logging.info('Start from epoch %s', start_epoch)

    # 如果配置中启用了 Weights & Biases (cfg.wandb.use)，尝试导入 wandb 库。如果导入失败，则抛出错误。
    if cfg.wandb.use:
        try:
            import wandb
        except:
            raise ImportError('WandB is not installed.')
        # 如果未指定 WandB 的名称，则使用 make_wandb_name() 生成一个名称。接着，初始化一个新的 WandB 运行并更新配置。
        if cfg.wandb.name == '':
            # todo 改动的时候需要修改
            wandb_name = make_wandb_name(cfg)
        else:
            wandb_name = cfg.wandb.name
        run = wandb.init(entity=cfg.wandb.entity, project=cfg.wandb.project,
                         name=wandb_name)
        run.config.update(cfg_to_dict(cfg))

    # num_splits：记录日志记录器的数量。
    # split_names：定义验证集和测试集的名称。
    # full_epoch_times：用于存储每个轮次的耗时。
    # perf：用于存储每个数据集的性能指标。
    num_splits = len(loggers)
    split_names = ['val', 'test']
    full_epoch_times = []
    perf = [[] for _ in range(num_splits)]

    # 使用 for 循环遍历每个轮次，从 start_epoch 到最大轮数 (cfg.optim.max_epoch)。
    for cur_epoch in range(start_epoch, cfg.optim.max_epoch):
        # 记录当前时间 start_time，用于计算本轮的耗时
        start_time = time.perf_counter()
        # 训练当前轮次：调用 train_epoch() 函数进行训练。
        train_epoch(loggers[0], loaders[0], model, optimizer, scheduler,
                    cfg.optim.batch_accumulation)
        # 记录训练性能：将当前轮次的性能指标写入第一个日志记录器（训练集），并保存到 perf 列表中。
        perf[0].append(loggers[0].write_epoch(cur_epoch))

        # 评估模型性能：如果当前轮次是评估轮次（由 is_eval_epoch(cur_epoch) 判断），则对每个数据集（除了训练集）进行评估，调用 eval_epoch() 并记录性能。
        if is_eval_epoch(cur_epoch):
            # 验证集、测试集
            for i in range(1, num_splits):
                eval_epoch(loggers[i], loaders[i], model,
                           split=split_names[i - 1])
                perf[i].append(loggers[i].write_epoch(cur_epoch))
        else:
            # 如果当前轮次不是评估轮次，则将上一个轮次的性能指标追加到列表中，以保持一致。
            for i in range(1, num_splits):
                perf[i].append(perf[i][-1])

        # 更新学习率：从验证性能中提取最近的性能指标，并根据当前的学习率调度策略更新调度器。
        val_perf = perf[1]
        if cfg.optim.scheduler == 'reduce_on_plateau':
            scheduler.step(val_perf[-1]['loss'])
        else:
            scheduler.step()

        # 记录轮次耗时：计算本轮的总耗时并将其添加到 full_epoch_times 列表中。
        full_epoch_times.append(time.perf_counter() - start_time)

        # Checkpoint with regular frequency (if enabled).
        # 检查点保存：如果启用检查点保存，并且当前轮次是检查点轮次，则调用 save_ckpt() 函数保存模型、优化器和调度器的状态。
        if cfg.train.enable_ckpt and not cfg.train.ckpt_best \
                and is_ckpt_epoch(cur_epoch):
            save_ckpt(model, optimizer, scheduler, cur_epoch)

        # 记录 WandB 性能：如果启用了 WandB，则记录当前的性能指标。
        if cfg.wandb.use:
            run.log(flatten_dict(perf), step=cur_epoch)

        # Log current best stats on eval epoch.

        # 评估当前最佳性能
        if is_eval_epoch(cur_epoch):
            # 如果当前轮次是评估轮次，找到验证性能中损失最小的轮次
            best_epoch = np.array([vp['loss'] for vp in val_perf]).argmin()
            # 初始化最佳训练、验证和测试性能字符串。
            best_train = best_val = best_test = ""
            # 自定义最佳指标：如果配置中指定了最佳指标（cfg.metric_best），则根据该指标从验证性能中找到最佳轮次。
            if cfg.metric_best != 'auto':
                # Select again based on val perf of `cfg.metric_best`.
                m = cfg.metric_best
                # 从 val_perf 中提取出每个轮次的指标 m，生成一个包含这些指标值的列表
                # cfg.metric_agg 是一个字符串，指定了你想要对这些指标值进行什么聚合操作（例如 argmin、argmax、mean 等）。该字符串对应的是 numpy 数组的一个方法。
                # 最后，调用由 getattr() 动态获取到的方法，完成聚合操作。
                # 例如，argmin() 方法会返回最小值的索引，argmax() 会返回最大值的索引。
                best_epoch = getattr(np.array([vp[m] for vp in val_perf]),
                                     cfg.metric_agg)()

                # 记录最佳训练性能：检查最佳轮次的训练性能，如果没有，则默认值为0。
                if m in perf[0][best_epoch]:
                    best_train = f"train_{m}: {perf[0][best_epoch][m]:.4f}"
                else:
                    # Note: For some datasets it is too expensive to compute
                    # the main metric on the training set.
                    best_train = f"train_{m}: {0:.4f}"
                # 记录最佳验证和测试性能：获取最佳轮次的验证和测试性能。
                best_val = f"val_{m}: {perf[1][best_epoch][m]:.4f}"
                best_test = f"test_{m}: {perf[2][best_epoch][m]:.4f}"

                # 记录最佳性能指标到 WandB：将最佳轮次的各项性能指标记录到 WandB，包括损失和其他自定义指标。
                if cfg.wandb.use:
                    bstats = {"best/epoch": best_epoch}
                    for i, s in enumerate(['train', 'val', 'test']):
                        bstats[f"best/{s}_loss"] = perf[i][best_epoch]['loss']
                        if m in perf[i][best_epoch]:
                            bstats[f"best/{s}_{m}"] = perf[i][best_epoch][m]
                            run.summary[f"best_{s}_perf"] = \
                                perf[i][best_epoch][m]
                        for x in ['hits@1', 'hits@3', 'hits@10', 'mrr']:
                            if x in perf[i][best_epoch]:
                                bstats[f"best/{s}_{x}"] = perf[i][best_epoch][x]
                    run.log(bstats, step=cur_epoch)
                    run.summary["full_epoch_time_avg"] = np.mean(full_epoch_times)
                    run.summary["full_epoch_time_sum"] = np.sum(full_epoch_times)

            # Checkpoint the best epoch params (if enabled).
            # 保存最佳轮次检查点：如果启用了检查点保存并且当前轮次是最佳轮次，则保存检查点。如果启用了清理，则调用 clean_ckpt() 删除旧的检查点。
            if cfg.train.enable_ckpt and cfg.train.ckpt_best and \
                    best_epoch == cur_epoch:
                save_ckpt(model, optimizer, scheduler, cur_epoch)
                if cfg.train.ckpt_clean:  # Delete old ckpt each time.
                    clean_ckpt()

            # 记录当前轮次信息：使用日志记录当前轮次的时间、平均时间、最佳轮次及其训练、验证和测试损失。
            logging.info(
                f"> Epoch {cur_epoch}: took {full_epoch_times[-1]:.1f}s "
                f"(avg {np.mean(full_epoch_times):.1f}s) | "
                f"Best so far: epoch {best_epoch}\t"
                f"train_loss: {perf[0][best_epoch]['loss']:.4f} {best_train}\t"
                f"val_loss: {perf[1][best_epoch]['loss']:.4f} {best_val}\t"
                f"test_loss: {perf[2][best_epoch]['loss']:.4f} {best_test}"
            )

            # 记录模型参数：如果模型包含 trf_layers 属性，遍历这些层并记录可训练的 gamma 参数值。
            if hasattr(model, 'trf_layers'):
                # Log SAN's gamma parameter values if they are trainable.
                for li, gtl in enumerate(model.trf_layers):
                    if torch.is_tensor(gtl.attention.gamma) and \
                            gtl.attention.gamma.requires_grad:
                        logging.info(f"    {gtl.__class__.__name__} {li}: "
                                     f"gamma={gtl.attention.gamma.item()}")

    # 记录平均耗时：在所有轮次完成后，记录每轮的平均时间和总时间。
    logging.info(f"Avg time per epoch: {np.mean(full_epoch_times):.2f}s")
    logging.info(f"Total train loop time: {np.sum(full_epoch_times) / 3600:.2f}h")
    # 关闭日志记录器：在训练结束后，关闭所有日志记录器。
    for logger in loggers:
        logger.close()
    # 清理检查点：如果配置中启用了清理，删除多余的检查点。
    if cfg.train.ckpt_clean:
        clean_ckpt()
    # close wandb
    # 结束 WandB 运行：如果使用了 WandB，结束当前的运行并清空引用。
    if cfg.wandb.use:
        run.finish()
        run = None

    # 记录任务完成：最后，记录任务完成的信息，包括结果保存的目录。
    logging.info('Task done, results saved in %s', cfg.run_dir)


# inference_only() 函数的作用是运行一个推理流程，在不进行模型训练或更新的情况下，评估模型在验证集或测试集上的性能。
@register_train('inference-only')
def inference_only(loggers, loaders, model, optimizer=None, scheduler=None):
    """
    Customized pipeline to run inference only.

    Args:
        loggers: List of loggers
        loaders: List of loaders
        model: GNN model
        optimizer: Unused, exists just for API compatibility
        scheduler: Unused, exists just for API compatibility
    """
    num_splits = len(loggers)
    split_names = ['train', 'val', 'test']
    # 初始化一个性能记录列表，perf 列表的每个子列表对应一个数据集（如验证集和测试集）。每次评估后，会将结果追加到对应的数据集子列表中。
    perf = [[] for _ in range(num_splits)]
    # 因为推理没有训练过程，这里的 cur_epoch 固定为 0。
    cur_epoch = 0
    # 记录推理开始的时间，用于计算总耗时。
    start_time = time.perf_counter()

    # 遍历每个数据集（训练集、验证集、测试集）并进行推理。
    for i in range(0, num_splits):
        eval_epoch(loggers[i], loaders[i], model,
                   split=split_names[i])
        # 将该 epoch 的推理结果记录到 perf[i] 中，并通过 loggers[i].write_epoch(cur_epoch) 写入日志文件。由于推理没有迭代，因此 epoch 固定为 0。
        perf[i].append(loggers[i].write_epoch(cur_epoch))
    # 获取验证集的推理结果，通常验证集用来判断模型的表现。
    val_perf = perf[1]

    # 计算验证集中损失值最小的 epoch，找出表现最好的 epoch。
    best_epoch = np.array([vp['loss'] for vp in val_perf]).argmin()
    best_train = best_val = best_test = ""
    # 如果配置文件中 cfg.metric_best 指定了一个非默认的指标（不是 auto），则使用该指标来选择最优性能。
    if cfg.metric_best != 'auto':
        # Select again based on val perf of `cfg.metric_best`.
        # 获取用户指定的性能指标（如 accuracy、f1）。
        m = cfg.metric_best
        # 根据验证集中指定指标 m，通过 cfg.metric_agg（如 min 或 max）找到最优 epoch。
        best_epoch = getattr(np.array([vp[m] for vp in val_perf]),
                             cfg.metric_agg)()
        # 分别存储在最优 epoch 时训练集、验证集和测试集的指定性能指标。
        if m in perf[0][best_epoch]:
            best_train = f"train_{m}: {perf[0][best_epoch][m]:.4f}"
        else:
            # Note: For some datasets it is too expensive to compute
            # the main metric on the training set.
            best_train = f"train_{m}: {0:.4f}"
        best_val = f"val_{m}: {perf[1][best_epoch][m]:.4f}"
        best_test = f"test_{m}: {perf[2][best_epoch][m]:.4f}"

    # 使用日志记录推理过程的总结，输出最优 epoch 时训练集、验证集、测试集的损失值和其他指定指标。
    logging.info(
        f"> Inference | "
        f"train_loss: {perf[0][best_epoch]['loss']:.4f} {best_train}\t"
        f"val_loss: {perf[1][best_epoch]['loss']:.4f} {best_val}\t"
        f"test_loss: {perf[2][best_epoch]['loss']:.4f} {best_test}"
    )
    # 记录推理任务完成的总耗时。
    logging.info(f'Done! took: {time.perf_counter() - start_time:.2f}s')
    # 关闭每个 logger，确保日志数据写入文件。
    for logger in loggers:
        logger.close()
